{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'action' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[9], line 78\u001b[0m\n\u001b[0;32m     75\u001b[0m state \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mexpand_dims(state, axis\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m)  \u001b[38;5;66;03m# 添加 batch_size 维度\u001b[39;00m\n\u001b[0;32m     76\u001b[0m state \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mtranspose(state, (\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m3\u001b[39m, \u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m2\u001b[39m))  \u001b[38;5;66;03m# 转换为 [batch_size, channels, height, width]\u001b[39;00m\n\u001b[1;32m---> 78\u001b[0m next_state, reward, done, _ \u001b[38;5;241m=\u001b[39m env\u001b[38;5;241m.\u001b[39mstep(\u001b[43maction\u001b[49m)\n\u001b[0;32m     79\u001b[0m next_state \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mexpand_dims(next_state, axis\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m)\n\u001b[0;32m     80\u001b[0m next_state \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mtranspose(next_state, (\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m3\u001b[39m, \u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m2\u001b[39m))\n",
      "\u001b[1;31mNameError\u001b[0m: name 'action' is not defined"
     ]
    }
   ],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "import gym\n",
    "from gym import spaces\n",
    "\n",
    "class SnakeGame(gym.Env):\n",
    "    def __init__(self):\n",
    "        super(SnakeGame, self).__init__()\n",
    "        \n",
    "        self.grid_size = 10  # 游戏网格大小\n",
    "        self.snake = [(5, 5)]  # 初始蛇的位置\n",
    "        self.food = (random.randint(0, self.grid_size - 1), random.randint(0, self.grid_size - 1))  # 食物的位置\n",
    "        self.direction = (0, 1)  # 初始方向（向右）\n",
    "        self.done = False\n",
    "        self.reward = 0\n",
    "        \n",
    "        # 定义动作空间和状态空间\n",
    "        self.action_space = spaces.Discrete(4)  # 上，下，左，右\n",
    "        self.observation_space = spaces.Box(low=0, high=self.grid_size-1, shape=(2, self.grid_size, self.grid_size), dtype=np.int32)\n",
    "    \n",
    "    def step(self, action):\n",
    "        if action == 0:  # 上\n",
    "            self.direction = (-1, 0)\n",
    "        elif action == 1:  # 下\n",
    "            self.direction = (1, 0)\n",
    "        elif action == 2:  # 左\n",
    "            self.direction = (0, -1)\n",
    "        elif action == 3:  # 右\n",
    "            self.direction = (0, 1)\n",
    "\n",
    "        head_x, head_y = self.snake[0]\n",
    "        new_head = (head_x + self.direction[0], head_y + self.direction[1])\n",
    "\n",
    "        # 碰到墙壁或者自己\n",
    "        if (new_head[0] < 0 or new_head[1] < 0 or new_head[0] >= self.grid_size or new_head[1] >= self.grid_size or new_head in self.snake):\n",
    "            self.done = True\n",
    "            self.reward = -1  # 撞墙或撞到自己\n",
    "            return self._get_obs(), self.reward, self.done, {}\n",
    "        \n",
    "        # 吃到食物\n",
    "        if new_head == self.food:\n",
    "            self.snake.insert(0, new_head)\n",
    "            self.food = (random.randint(0, self.grid_size - 1), random.randint(0, self.grid_size - 1))  # 随机生成新的食物\n",
    "            self.reward = 1  # 吃到食物\n",
    "        else:\n",
    "            self.snake = [new_head] + self.snake[:-1]  # 移动蛇\n",
    "            self.reward = 0  # 无奖励\n",
    "\n",
    "        return self._get_obs(), self.reward, self.done, {}\n",
    "\n",
    "    def reset(self):\n",
    "        self.snake = [(5, 5)]  # 初始蛇的位置\n",
    "        self.food = (random.randint(0, self.grid_size - 1), random.randint(0, self.grid_size - 1))  # 随机生成食物\n",
    "        self.direction = (0, 1)  # 初始方向（向右）\n",
    "        self.done = False\n",
    "        return self._get_obs()\n",
    "\n",
    "    def render(self):\n",
    "        grid = np.zeros((self.grid_size, self.grid_size), dtype=int)\n",
    "        for segment in self.snake:\n",
    "            grid[segment[0], segment[1]] = 1  # 蛇身\n",
    "        grid[self.food[0], self.food[1]] = 2  # 食物\n",
    "        print(grid)\n",
    "        \n",
    "    # 修改 _get_obs 方法，返回正确的形状\n",
    "    def _get_obs(self):\n",
    "        grid = np.zeros((2, self.grid_size, self.grid_size), dtype=int)\n",
    "        for segment in self.snake:\n",
    "            grid[0, segment[0], segment[1]] = 1  # 蛇身\n",
    "        grid[1, self.food[0], self.food[1]] = 1  # 食物\n",
    "        return grid\n",
    "\n",
    "# 修改训练代码中的 state 和 next_state 维度\n",
    "state = env.reset()\n",
    "state = np.expand_dims(state, axis=0)  # 添加 batch_size 维度\n",
    "state = np.transpose(state, (0, 3, 1, 2))  # 转换为 [batch_size, channels, height, width]\n",
    "\n",
    "next_state, reward, done, _ = env.step(action)\n",
    "next_state = np.expand_dims(next_state, axis=0)\n",
    "next_state = np.transpose(next_state, (0, 3, 1, 2))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DQN(nn.Module):\n",
    "    def __init__(self, input_shape, n_actions):\n",
    "        super(DQN, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(input_shape[0], 32, kernel_size=3, stride=1)\n",
    "        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1)\n",
    "        self.fc1 = nn.Linear(64 * 6 * 6, 512)\n",
    "        self.fc2 = nn.Linear(512, n_actions)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = torch.relu(self.conv1(x))\n",
    "        x = torch.relu(self.conv2(x))\n",
    "        x = x.view(x.size(0), -1)  # Flatten\n",
    "        x = torch.relu(self.fc1(x))\n",
    "        return self.fc2(x)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Agent:\n",
    "    def __init__(self, input_shape, n_actions):\n",
    "        self.input_shape = input_shape\n",
    "        self.n_actions = n_actions\n",
    "        self.policy_net = DQN(input_shape, n_actions).to(device)\n",
    "        self.target_net = DQN(input_shape, n_actions).to(device)\n",
    "        self.target_net.load_state_dict(self.policy_net.state_dict())\n",
    "        self.target_net.eval()\n",
    "        \n",
    "        self.optimizer = optim.Adam(self.policy_net.parameters(), lr=0.0001)\n",
    "        self.memory = deque(maxlen=10000)\n",
    "        self.batch_size = 32\n",
    "        self.gamma = 0.99\n",
    "        self.epsilon = 0.1\n",
    "        self.epsilon_decay = 0.995\n",
    "        self.epsilon_min = 0.01\n",
    "\n",
    "    def select_action(self, state):\n",
    "        if random.random() < self.epsilon:\n",
    "            return random.randint(0, self.n_actions - 1)  # 探索\n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                state = torch.tensor(state, dtype=torch.float32).unsqueeze(0).to(device)\n",
    "                q_values = self.policy_net(state)\n",
    "                return torch.argmax(q_values).item()  # 利用\n",
    "\n",
    "    def store_transition(self, state, action, reward, next_state, done):\n",
    "        self.memory.append((state, action, reward, next_state, done))\n",
    "\n",
    "    def train(self):\n",
    "        if len(self.memory) < self.batch_size:\n",
    "            return\n",
    "        \n",
    "        batch = random.sample(self.memory, self.batch_size)\n",
    "        states, actions, rewards, next_states, dones = zip(*batch)\n",
    "        \n",
    "        states = torch.tensor(states, dtype=torch.float32).to(device)\n",
    "        next_states = torch.tensor(next_states, dtype=torch.float32).to(device)\n",
    "\n",
    "        actions = torch.tensor(actions).to(device)\n",
    "        rewards = torch.tensor(rewards).to(device)\n",
    "        dones = torch.tensor(dones).to(device)\n",
    "        \n",
    "        q_values = self.policy_net(states)\n",
    "        next_q_values = self.target_net(next_states)\n",
    "        \n",
    "        target_q_values = rewards + (1 - dones) * self.gamma * next_q_values.max(1)[0]\n",
    "        \n",
    "        loss = nn.MSELoss()(q_values.gather(1, actions.unsqueeze(1)), target_q_values.unsqueeze(1))\n",
    "        \n",
    "        self.optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        self.optimizer.step()\n",
    "        \n",
    "        self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)\n",
    "\n",
    "    def update_target_net(self):\n",
    "        self.target_net.load_state_dict(self.policy_net.state_dict())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "axes don't match array",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[8], line 14\u001b[0m\n\u001b[0;32m     12\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m episode \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(num_episodes):\n\u001b[0;32m     13\u001b[0m     state \u001b[38;5;241m=\u001b[39m env\u001b[38;5;241m.\u001b[39mreset()\n\u001b[1;32m---> 14\u001b[0m     state \u001b[38;5;241m=\u001b[39m \u001b[43mstate\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtranspose\u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m3\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# 调整为 [batch_size, height, width, channels]\u001b[39;00m\n\u001b[0;32m     15\u001b[0m     done \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[0;32m     16\u001b[0m     total_reward \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n",
      "\u001b[1;31mValueError\u001b[0m: axes don't match array"
     ]
    }
   ],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# 创建游戏环境\n",
    "env = SnakeGame()\n",
    "input_shape = (2, env.grid_size, env.grid_size)\n",
    "n_actions = env.action_space.n\n",
    "\n",
    "# 创建智能体\n",
    "agent = Agent(input_shape, n_actions)\n",
    "\n",
    "num_episodes = 1000\n",
    "for episode in range(num_episodes):\n",
    "    state = env.reset()\n",
    "    state = state.transpose((0, 2, 3, 1))  # 调整为 [batch_size, height, width, channels]\n",
    "    done = False\n",
    "    total_reward = 0\n",
    "    \n",
    "    while not done:\n",
    "        action = agent.select_action(state)\n",
    "        next_state, reward, done, _ = env.step(action)\n",
    "        next_state = next_state.transpose((0, 2, 3, 1))\n",
    "        \n",
    "        agent.store_transition(state, action, reward, next_state, done)\n",
    "        agent.train()\n",
    "        \n",
    "        state = next_state\n",
    "        total_reward += reward\n",
    "        \n",
    "        env.render()\n",
    "\n",
    "    print(f\"Episode {episode + 1}/{num_episodes}, Total Reward: {total_reward}\")\n",
    "\n",
    "    # 每50轮更新一次目标网络\n",
    "    if (episode + 1) % 50 == 0:\n",
    "        agent.update_target_net()\n",
    "\n",
    "env.close()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
